import sqlite3
import os
from datetime import datetime

# 数据库配置
DATABASE = 'todos.db'

def migrate_db():
    """迁移数据库，添加token和token_expiry列到users表"""
    print("开始迁移数据库...")
    
    # 检查数据库文件是否存在
    if not os.path.exists(DATABASE):
        print(f"数据库文件 {DATABASE} 不存在，无需迁移")
        return
    
    # 连接数据库
    conn = sqlite3.connect(DATABASE)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    
    # 检查users表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='users'")
    if cursor.fetchone() is None:
        print("users表不存在，无需迁移")
        conn.close()
        return
    
    # 检查token列是否已存在
    cursor.execute("PRAGMA table_info(users)")
    columns = cursor.fetchall()
    has_token = any(col['name'] == 'token' for col in columns)
    has_token_expiry = any(col['name'] == 'token_expiry' for col in columns)
    
    try:
        if not has_token:
            print("添加token列到users表...")
            cursor.execute("ALTER TABLE users ADD COLUMN token TEXT")
        
        if not has_token_expiry:
            print("添加token_expiry列到users表...")
            cursor.execute("ALTER TABLE users ADD COLUMN token_expiry TIMESTAMP")
        
        conn.commit()
        print("迁移成功完成！")
    except Exception as e:
        conn.rollback()
        print(f"迁移失败: {e}")
    finally:
        conn.close()

if __name__ == "__main__":
    migrate_db()
